Goto

Collaborating Authors

 healpix mesh


HEAL-ViT: Vision Transformers on a spherical mesh for medium-range weather forecasting

Ramavajjala, Vivek

arXiv.org Artificial Intelligence

In recent years, a variety of ML architectures and techniques have seen success in producing skillful medium range weather forecasts. In particular, Vision Transformer (ViT)-based models (e.g. Pangu-Weather, FuXi) have shown strong performance, working nearly "out-of-the-box" by treating weather data as a multi-channel image on a rectilinear grid. While a rectilinear grid is appropriate for 2D images, weather data is inherently spherical and thus heavily distorted at the poles on a rectilinear grid, leading to disproportionate compute being used to model data near the poles. Graph-based methods (e.g. GraphCast) do not suffer from this problem, as they map the longitude-latitude grid to a spherical mesh, but are generally more memory intensive and tend to need more compute resources for training and inference. While spatially homogeneous, the spherical mesh does not lend itself readily to be modeled by ViT-based models that implicitly rely on the rectilinear grid structure. We present HEAL-ViT, a novel architecture that uses ViT models on a spherical mesh, thus benefiting from both the spatial homogeneity enjoyed by graph-based models and efficient attention-based mechanisms exploited by transformers. HEAL-ViT produces weather forecasts that outperform the ECMWF IFS on key metrics, and demonstrate better bias accumulation and blurring than other ML weather prediction models. Further, the lowered compute footprint of HEAL-ViT makes it attractive for operational use as well, where other models in addition to a 6-hourly prediction model may be needed to produce the full set of operational forecasts required.


Advancing Parsimonious Deep Learning Weather Prediction using the HEALPix Mesh

Karlbauer, Matthias, Cresswell-Clay, Nathaniel, Moreno, Raul A., Durran, Dale R., Kurth, Thorsten, Butz, Martin V.

arXiv.org Artificial Intelligence

We present a parsimonious deep learning weather prediction model on the Hierarchical Equal Area isoLatitude Pixelization (HEALPix) to forecast seven atmospheric variables for arbitrarily long lead times on a global approximately 110 km mesh at 3h time resolution. In comparison to state-of-the-art machine learning weather forecast models, such as Pangu-Weather and GraphCast, our DLWP-HPX model uses coarser resolution and far fewer prognostic variables. Yet, at one-week lead times its skill is only about one day behind the state-of-the-art numerical weather prediction model from the European Centre for Medium-Range Weather Forecasts. We report successive forecast improvements resulting from model design and data-related decisions, such as switching from the cubed sphere to the HEALPix mesh, inverting the channel depth of the U-Net, and introducing gated recurrent units (GRU) on each level of the U-Net hierarchy. The consistent east-west orientation of all cells on the HEALPix mesh facilitates the development of location-invariant convolution kernels that are successfully applied to propagate global weather patterns across our planet. Without any loss of spectral power after two days, the model can be unrolled autoregressively for hundreds of steps into the future to generate stable and realistic states of the atmosphere that respect seasonal trends, as showcased in one-year simulations. Our parsimonious DLWP-HPX model is research-friendly and potentially well-suited for sub-seasonal and seasonal forecasting.